{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import input_data\n",
    "import fsl_model as tm\n",
    "import train\n",
    "from torch.utils.data import random_split\n",
    "from train import MyDataset, DataLoader\n",
    "args = {\n",
    "    # word embedding parameters\n",
    "    'vocab_size': -1, # vocab size (emb.shape[0])\n",
    "    'word_emb_dim': -1,# embed dim (emb.shape[1])\n",
    "    'hidden_dim': 32, # hidden emb dim\n",
    "    # attention parameters\n",
    "    'seq_len': -1, # max_time in original model, max len of all\n",
    "    'head_num': 6, # head number\n",
    "    'atte_dim': 20, # dim of attention\n",
    "    # capsnet parameters\n",
    "    's_cnum': -1, # seen class num\n",
    "    'u_cnum': -1, # unseen class num\n",
    "    'keep_prob': 0.8, # capnet dropout keep rate\n",
    "    'routing': 4, # capsule routing iter num\n",
    "    'out_atoms': -1, # equals to seen class num\n",
    "    'out_dim': 10, # capsule output dim\n",
    "    # learning parameters\n",
    "    'margin': 1, # ranking loss margin\n",
    "    'batch_size': 64,\n",
    "    'learning_rate': 1e-4\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------------------read datasets begin-------------------\n",
      "------------------load word2vec begin-------------------\n",
      "------------------load word2vec end---------------------\n",
      "------------------read datasets end---------------------\n"
     ]
    }
   ],
   "source": [
    "data = input_data.read_datasets()\n",
    "train_set = MyDataset(data['x_tr'], data['s_len'], data['y_tr'])\n",
    "tr_size = len(train_set) * 4 // 5\n",
    "va_size = len(train_set) - tr_size\n",
    "train_dataset, valid_dataset = random_split(train_set, [tr_size, va_size])\n",
    "# 划分训练、验证集\n",
    "train_loader = DataLoader(train_dataset, batch_size=args['batch_size'], shuffle=True, num_workers=2)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=args['batch_size'], shuffle=True, num_workers=2)\n",
    "# 训练模型\n",
    "model = tm.FSL_model(data, args)\n",
    "train.train(model, train_loader, epoch=20, lr=1e-4, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Val Accuracy: 97.02%\n"
     ]
    }
   ],
   "source": [
    "# 进行验证\n",
    "train.validation(model, valid_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 80.36%\n"
     ]
    }
   ],
   "source": [
    "# FSL预测\n",
    "fsl_set = MyDataset(data['x_te'], data['u_len'], data['y_te'])\n",
    "fsl_loader = DataLoader(fsl_set, batch_size=args['batch_size'], shuffle=True, num_workers=2)\n",
    "train.FSL_testing(model, fsl_loader)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "a00b85c57b695ce9c151329fc92e9a120e775a6772755d2fc81cbf2708e2f678"
  },
  "kernelspec": {
   "display_name": "Python 3.7.10 64-bit ('cs224n': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}